import cv2 as cv
from ultralytics import YOLO
import math
import time

# 0:鼻子1:左眼2:右眼3:左耳4:右耳5:左肩6:右肩7:左肘8:右肘9:左腕10:右腕11:左髋12:右髋13:左膝14:右膝15:左踝16:右踝

# 加载模型
pose_model = YOLO('yolov8n-pose.pt')  # 姿态估计模型
phone_model = YOLO('yolov8n.pt')  # 手机检测模型
# cap = cv2.VideoCapture("http://192.168.188.250:4747/mjbegfeed") # 手机接入“我高贵的校园网回来啦”移动热点时的视像头url
cap = cv.VideoCapture(0)
epoch = 0

while True:
    ret, frame = cap.read()
    if not ret:
        # 如果读取帧失败（例如，视频结束），则退出循环
        break
        # 检测手机
    phone_result = phone_model(frame)
    phone_result = phone_result[0]
    phone_boxes = phone_result.boxes
    phone_detected = False
    class_ids = phone_boxes.cls.cpu().numpy().astype(int)
    if 67 in class_ids:  # 检测到手机
        for phone_box in phone_boxes:
            xyxy_tensor = phone_box.xyxy[0]  # 提取cellphone边界框的xyxy信息
            xyxy_coords = xyxy_tensor.cpu().numpy()
            x1, y1, x2, y2 = xyxy_coords
        phone_detected = True

    # 检测姿态
    if phone_detected:
        pose_results = pose_model(frame)
        results = pose_model(frame)
        keypoints = results[0].keypoints

        xy_cpu = keypoints.xy.cpu() if keypoints.xy.device != 'cpu' else keypoints.xy

        # 根据官方模型的说明书，从17个点位中提取左右臂共6个点位
        shoulder_x_r, shoulder_y_r = xy_cpu[0, 6, 0].item(), xy_cpu[0, 6, 1].item()
        elbow_x_r, elbow_y_r = xy_cpu[0, 8, 0].item(), xy_cpu[0, 8, 1].item()
        wrist_x_r, wrist_y_r = xy_cpu[0, 10, 0].item(), xy_cpu[0, 10, 1].item()
        shoulder_x_l, shoulder_y_l = xy_cpu[0, 5, 0].item(), xy_cpu[0, 5, 1].item()
        elbow_x_l, elbow_y_l = xy_cpu[0, 7, 0].item(), xy_cpu[0, 7, 1].item()
        wrist_x_l, wrist_y_l = xy_cpu[0, 9, 0].item(), xy_cpu[0, 9, 1].item()

        # 根据官方模型，从17个点位提取髋关节位置，通过判断肩至髋的长度与肩宽的关系，判断是否处于侧身
        hip_y_l = xy_cpu[0, 11, 1]
        hip_y_r = xy_cpu[0, 12, 1]

        # 计算左右臂的四个向量
        vector_A_x_r = elbow_x_r - shoulder_x_r
        vector_A_y_r = elbow_y_r - shoulder_y_r

        vector_B_x_r = wrist_x_r - elbow_x_r
        vector_B_y_r = wrist_y_r - elbow_y_r

        vector_A_x_l = elbow_x_l - shoulder_x_l
        vector_A_y_l = elbow_y_l - shoulder_y_l

        vector_B_x_l = wrist_x_l - elbow_x_l
        vector_B_y_l = wrist_y_l - elbow_y_l

        # 计算左右臂的大臂、小臂向量乘积
        dot_product_r = vector_A_x_r * vector_B_x_r + vector_A_y_r * vector_B_y_r
        dot_product_l = vector_A_x_l * vector_B_x_l + vector_A_y_l * vector_B_y_l

        # 计算长度
        magnitude_A_r = math.sqrt(vector_A_x_r ** 2 + vector_A_y_r ** 2)
        magnitude_B_r = math.sqrt(vector_B_x_r ** 2 + vector_B_y_r ** 2)
        magnitude_A_l = math.sqrt(vector_A_x_l ** 2 + vector_A_y_l ** 2)
        magnitude_B_l = math.sqrt(vector_B_x_l ** 2 + vector_B_y_l ** 2)

        # 保障计算时分母不为0
        if magnitude_A_r * magnitude_B_r == 0:
            continue
        if magnitude_A_l * magnitude_B_l == 0:
            continue

        # 计算cos值
        cosine_angle_r = dot_product_r / (magnitude_A_r * magnitude_B_r)
        cosine_angle_l = dot_product_l / (magnitude_A_l * magnitude_B_l)

        # 计算角度
        # angle_in_degrees_r = math.degrees(math.acos(cosine_angle_r))
        # angle_in_degrees_l = math.degrees(math.acos(cosine_angle_l))

        # 找出最接近九十度的夹角
        # closest_angle = min(angle_in_degrees_r, angle_in_degrees_l, key=lambda x: abs(x - 98.12))
        closest_angle = min(cosine_angle_l, cosine_angle_r)

        # 首先计算手机的中心位置
        x_phone_mid = (x1 + x2) / 2
        y_phone_mid = (y1 + y2) / 2

        # magnitude_B_l和magnitude_B_r为左右侧的小臂长度
        # 计算手机中心到两手腕的距离，取最小的一个作为所需
        phone_to_leftwrist = math.sqrt((x_phone_mid - wrist_x_l) ** 2 + (y_phone_mid - wrist_y_l) ** 2)
        phone_to_rightwrist = math.sqrt((x_phone_mid - wrist_x_r) ** 2 + (y_phone_mid - wrist_y_r) ** 2)
        phone_to_wirst = min(phone_to_rightwrist, phone_to_leftwrist)

        # 判断是否为侧身：肩宽/髋宽
        height_l = shoulder_y_l - hip_y_l
        height_r = shoulder_y_r - hip_y_r
        delta_y_hiptoshoulder = max(height_l, height_r)
        delta_x_shoulder = math.sqrt((shoulder_x_r - shoulder_x_l) ** 2 + (shoulder_y_r - shoulder_y_l) ** 2)
        bili = delta_x_shoulder / delta_y_hiptoshoulder

        # 检查角度是否在指定范围内，并且手机被检测到，且是该人在走路玩手机
        if -0.5345 <= closest_angle <= 0.2756 and phone_to_wirst <= 0.4 * magnitude_B_l:
            # 因为存不到照片，减弱判断范围
            if 0.09080594 <= abs(bili) <= 0.36149454:
                cv.imwrite('./images/result_' + str(epoch) + '.jpg', frame)
                epoch += 1
                print("检测到走路玩手机的同学，已经保存")

    # 显示图像
    cv.imshow("image", frame)
    k = cv.waitKey(1)
    time.sleep(1000)
    if k == ord("q"):
        break

cap.release()
cv.destroyAllWindows()
